package com.bbm.communitysystem.core.ai;

import com.bbm.communitysystem.domain.dto.chat.ChatHistoryDto;
import com.bbm.communitysystem.domain.vo.ChatVo;
import com.bbm.communitysystem.service.impl.ChatHistoryServiceImpl;
import com.bbm.communitysystem.tools.DemoDate;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.concurrent.atomic.AtomicReference;

@Component
@RequiredArgsConstructor
public class Ai {

    private final OllamaChatModel ollamaChatModel;

    private final ChatHistoryServiceImpl chatHistoryService;



    public Flux<ChatVo> chat(Prompt prompt, String userId, Long pageId, String note) {

        String finalContent = prompt.getContents();

        AtomicReference<Boolean> isThink = new AtomicReference<>(false);
        StringBuilder con = new StringBuilder();
        StringBuilder thinkContent = new StringBuilder();

        Flux<ChatVo> map = ollamaChatModel.stream(prompt).map(chatResponse -> {

            ChatVo chatVo = new ChatVo();
            String text1 = chatResponse.getResult().getOutput().getText();

            if (text1.equals("<think>")) {
                isThink.set(true);
            }

            if (text1.equals("</think>")) {
                isThink.set(false);
            }

            if (isThink.get() || text1.equals("</think>") || text1.equals("<think>")) {
                thinkContent.append(text1);
            } else {
                con.append(text1);
            }

            chatVo.setIsThink(isThink.get() ? 1 : 0);
            chatVo.setShowTime(DemoDate.getPostDate(System.currentTimeMillis()));
            chatVo.setEvent("message");
            chatVo.setIsShow("1");
            chatVo.setData("none");
            chatVo.setType("text");
            chatVo.setContent(text1);

            if (text1.isEmpty()) {
                chatVo.setEvent("end");
                StoreChatHistory(userId, note, con.toString(), thinkContent.toString(), pageId);
            }



            return chatVo;
        });




        return map;
    }

    public Flux<ChatVo> chat(Prompt prompt, String userId, Long pageId, Boolean isStore) {

        if (isStore) {
            return chat(prompt, userId, pageId);
        } else {
            return chat(prompt);
        }
    }

    public Flux<ChatVo> chat(Prompt prompt) {
        String finalContent = prompt.getContents();

        AtomicReference<Boolean> isThink = new AtomicReference<>(false);

        Flux<ChatVo> map = ollamaChatModel.stream(prompt).map(chatResponse -> {

            ChatVo chatVo = new ChatVo();
            String text1 = chatResponse.getResult().getOutput().getText();

            if (text1.equals("<think>")) {
                isThink.set(true);
            }

            if (text1.equals("</think>")) {
                isThink.set(false);
            }


            chatVo.setIsThink(isThink.get() ? 1 : 0);
            chatVo.setShowTime(DemoDate.getPostDate(System.currentTimeMillis()));
            chatVo.setEvent("message");
            chatVo.setIsShow("1");
            chatVo.setData("none");
            chatVo.setType("text");
            chatVo.setContent(text1);

            if (text1.isEmpty()) {
                chatVo.setEvent("end");
            }

            return chatVo;
        });




        return map;
    }

    public Flux<ChatVo> chat(Prompt prompt, String userId, Long pageId) {

        String finalContent = prompt.getContents();

        AtomicReference<Boolean> isThink = new AtomicReference<>(false);
        StringBuilder con = new StringBuilder();
        StringBuilder thinkContent = new StringBuilder();

        Flux<ChatVo> map = ollamaChatModel.stream(prompt).map(chatResponse -> {

            ChatVo chatVo = new ChatVo();
            String text1 = chatResponse.getResult().getOutput().getText();

            if (text1.equals("<think>")) {
                isThink.set(true);
            }

            if (text1.equals("</think>")) {
                isThink.set(false);
            }

            if (isThink.get() || text1.equals("</think>") || text1.equals("<think>")) {
                thinkContent.append(text1);
            } else {
                con.append(text1);
            }

            chatVo.setIsThink(isThink.get() ? 1 : 0);
            chatVo.setShowTime(DemoDate.getPostDate(System.currentTimeMillis()));
            chatVo.setEvent("message");
            chatVo.setIsShow("1");
            chatVo.setData("none");
            chatVo.setType("text");
            chatVo.setContent(text1);

            if (text1.isEmpty()) {
                chatVo.setEvent("end");
                StoreChatHistory(userId, finalContent, con.toString(), thinkContent.toString(), pageId);
            }



            return chatVo;
        });




        return map;
    }

    public void StoreChatHistory(String userId, String content, String aiReply, String aiThinkContent, Long pageId) {
        ChatHistoryDto user = new ChatHistoryDto();
        // 保存到数据库
        user.setChatPageId(pageId);
        user.setContent(content);
        user.setType("user");
        user.setThinkContent("");
        user.setCreateTime(System.currentTimeMillis() + "");
        user.setShowTime(user.getCreateTime());
        user.setStatus(1);

        chatHistoryService.addChatHistory(user);

        ChatHistoryDto chatHistoryDto = new ChatHistoryDto();
        // 保存到数据库
        chatHistoryDto.setChatPageId(pageId);
        chatHistoryDto.setContent(aiReply);
        chatHistoryDto.setType("ai");
        chatHistoryDto.setThinkContent(aiThinkContent);
        chatHistoryDto.setCreateTime(System.currentTimeMillis() + "");
        chatHistoryDto.setShowTime(chatHistoryDto.getCreateTime());
        chatHistoryDto.setStatus(1);

        chatHistoryService.addChatHistory(chatHistoryDto);

    }

}
